import numpy as np
from FNN import *
from utils import *
from readdata import *
import matplotlib as mpl
import matplotlib.pyplot as plt


# _shape like [feature dim, num, num, ..., out dim]
# _activation='sigmoid' or 'softmax'
def train(_dataset, _shape, _lr=0.1, _activation='sigmoid', _epoch=20, _batchsize=1, _normalize=False, _plot=False):
    cost_func = mean_squared_cost
    if _activation == 'softmax':
        cost_func = cross_entropy_cost
    fnn = FNN(_shape, activation=_activation)
    x_train, label_train, x_test, label_test = read_data(_dataset)
    if _normalize:
        m, v = get_mean_variance(x_train)
        fnn.set_normalize(m, v)
    y_test = one_hot(_shape[-1], label_test)
    train_size = x_train.shape[0]
    d_lr = _lr
    max_acc = 0

    # plot
    N, M = 1000, 1000
    x1_min, x2_min = x_test[:, 0].min() - 0.5, x_test[:, 1].min() - 0.5
    x1_max, x2_max = x_test[:, 0].max() + 0.5, x_test[:, 1].max() + 0.5
    t1 = np.linspace(x1_min, x1_max, N)
    t2 = np.linspace(x2_min, x2_max, M)
    x1, x2 = np.meshgrid(t1, t2)
    x_show = np.stack((x1.flat, x2.flat), axis=1)
    if _shape[-1] == 3:
        cm_light = mpl.colors.ListedColormap(['#FFA0A0', '#A0FFA0', '#A0A0FF'])
        cm_dark = mpl.colors.ListedColormap(['g', 'r', 'b'])
    elif _shape[-1] == 2:
        cm_light = mpl.colors.ListedColormap(['#FFA0A0', '#A0FFA0'])
        cm_dark = mpl.colors.ListedColormap(['r', 'g'])
    plt.ion()

    for ep in range(_epoch):
        shuffle(x_train, label_train)
        y_train = one_hot(_shape[-1], label_train)
        for i in range(0, train_size, _batchsize):
            x_batch = x_train[i:i + _batchsize]
            y_batch = y_train[i:i + _batchsize]
            res = fnn.forward(x_batch)
            loss = cost_func(y_batch, res)
            fnn.backward(y_batch, lr=d_lr)
            print('\rloss : {}'.format(loss), end='', flush=True)
        pred = fnn.forward(x_test)
        loss = cost_func(y_test, pred)
        pred = np.argmax(pred, axis=1)
        assert pred.shape == label_test.shape
        acc = np.sum(pred == label_test) / label_test.shape[0]
        if acc > max_acc:
            max_acc = acc
            fnn.save(_dataset)
        d_lr = dynamic_lr(_lr, ep)
        print('\nepoch {} acc:{} loss: {}\n'.format(ep + 1, acc, loss))

        # plot
        if _plot:
            plt.clf()
            y_show = fnn.forward(x_show)
            y_show = np.argmax(y_show, axis=1)
            plt.pcolormesh(x1, x2, y_show.reshape(x1.shape),
                           shading='auto', cmap=cm_light)
            plt.scatter(x_test[:, 0], x_test[:, 1], c=label_test,
                        cmap=cm_dark, marker='o', edgecolors='k')
            plt.xlabel("feature_1")
            plt.ylabel("feature_2")
            plt.title("FNN on {} Epoch {}".format(_dataset, ep + 1))
            plt.grid(True, ls=':')
            plt.pause(0.1)
            plt.ioff()
    plt.show()
    return max_acc


# _dataset = 'Exam' or 'Iris'
# _shape like [feature dim, num, num, ..., out dim]
def test(_dataset, _shape):
    x_train, label_train, x_test, label_test = read_data(_dataset)
    fnn = FNN(_shape)
    fnn.load(_dataset + '.txt')
    pred = fnn.forward(x_test)
    pred = np.argmax(pred, axis=1)
    acc = np.sum(pred == label_test) / label_test.shape[0]
    print('acc = {}'.format(acc))


def main():
    out_dim = {'Iris': 3, 'Exam': 2}

    dataset = 'Exam'
    model = [2, 20, out_dim[dataset]]
    max_acc = train(dataset, model, _lr=0.01,
                    _activation='softmax', _normalize=True, _plot=True)
    print('Exam max acc is {}'.format(max_acc))

    # dataset = 'Iris'
    # model = [2, 20, out_dim[dataset]]
    # max_acc = train(dataset, model, _lr=0.01, _activation='softmax', _normalize=False, _plot=True)
    # print('Iris max acc is {}'.format(max_acc))


if __name__ == '__main__':
    main()
